#include "vkop_averagepool.h"
#include "vkinfer/core/utils.h"
#include "vkinfer/runtime/vk/vkutils.h"

namespace vkinfer
{
    VkOpAveragePool::VkOpAveragePool(VkDevice device, const std::shared_ptr<OnnxNodeInfo>& opinfo)
    {
        std::string shader_path = get_shader_file_path("averagepool.spv");
        uint32_t num_outputs = (uint32_t)opinfo->output_names.size();
        uint32_t num_inputs = (uint32_t)opinfo->input_names.size();
        uint32_t num_weights = (uint32_t)opinfo->weight_names.size();
        this->initialize(device, shader_path, num_outputs, num_inputs, num_weights, sizeof(pc));

        this->name = opinfo->name;
        this->input_names = opinfo->input_names;
        this->output_names = opinfo->output_names;
        this->type = OpType::OP_TYPE_AVERAGEPOOL;

        // initialize parameters;
        for (auto& attr : opinfo->attributes)
        {
            if (attr.name == "strides")
                pc.stride = attr.get_int(0);
            else if (attr.name == "kernel_shape")
                pc.kernel = attr.get_int(0);
            else if (attr.name == "pads")
                pc.pad = attr.get_int(0);
        }
    }

    void VkOpAveragePool::update_descriptors(
        const std::vector<std::shared_ptr<VkBlob>>& inputs,
        std::vector<std::shared_ptr<VkBlob>>& outputs)
    {
        // update descriptor set information.
        uint32_t num_descriptors = (uint32_t)inputs.size() + (uint32_t)outputs.size();
        std::vector<VkDescriptorBufferInfo> descriptorBufferInfos(num_descriptors);
        std::vector<VkWriteDescriptorSet> writeDescriptorSets(num_descriptors);

        int index = 0;
        // output bindings come first
        for (size_t i = 0; i < outputs.size(); ++i)
        {
            descriptorBufferInfos[index] = {};
            descriptorBufferInfos[index].buffer = outputs[i]->buffer;
            descriptorBufferInfos[index].offset = 0;
            descriptorBufferInfos[index].range = outputs[i]->size;

            writeDescriptorSets[index] = {};
            writeDescriptorSets[index].sType = VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET;
            writeDescriptorSets[index].dstSet = descriptorSet;
            writeDescriptorSets[index].dstBinding = index;
            writeDescriptorSets[index].descriptorCount = 1;
            writeDescriptorSets[index].descriptorType = VK_DESCRIPTOR_TYPE_STORAGE_BUFFER;
            writeDescriptorSets[index].pBufferInfo = &descriptorBufferInfos[index];
            index += 1;
        }
        // then input bindings
        for (size_t i = 0; i < inputs.size(); ++i)
        {
            descriptorBufferInfos[index] = {};
            descriptorBufferInfos[index].buffer = inputs[i]->buffer;
            descriptorBufferInfos[index].offset = 0;
            descriptorBufferInfos[index].range = inputs[i]->size;

            writeDescriptorSets[index] = {};
            writeDescriptorSets[index].sType = VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET;
            writeDescriptorSets[index].dstSet = descriptorSet;
            writeDescriptorSets[index].dstBinding = index;
            writeDescriptorSets[index].descriptorCount = 1;
            writeDescriptorSets[index].descriptorType = VK_DESCRIPTOR_TYPE_STORAGE_BUFFER;
            writeDescriptorSets[index].pBufferInfo = &descriptorBufferInfos[index];
            index += 1;
        }
        vkUpdateDescriptorSets(device, (uint32_t)writeDescriptorSets.size(), writeDescriptorSets.data(), 0, NULL);
    }

    void VkOpAveragePool::forward(const std::vector<std::shared_ptr<Blob>>& inputs,
        std::vector<std::shared_ptr<Blob>>& outputs)
    {
        // calculate output size
        std::vector<uint32_t> input_shape = inputs[0]->shape;
        uint32_t input_h = input_shape[2];
        uint32_t input_w = input_shape[3];
        uint32_t input_c = input_shape[1];
        uint32_t output_c = input_c;

        uint32_t output_w = (input_w + pc.pad - pc.kernel) / pc.stride + 1;
        uint32_t output_h = (input_h + pc.pad - pc.kernel) / pc.stride + 1;
        if (output_w <= 0) output_w = 1;
        if (output_h <= 0) output_h = 1;
        std::vector<uint32_t> output_shape = { input_shape[0], output_c, (uint32_t)output_h, (uint32_t)output_w };

        outputs[0]->resize(output_shape);

        // prepare io blobs
        std::vector<std::shared_ptr<VkBlob>> input_blobs;
        std::vector<std::shared_ptr<VkBlob>> output_blobs;
        for (auto b : inputs)
        {
            std::shared_ptr<VkBlob> vb = std::static_pointer_cast<VkBlob, Blob>(b);
            input_blobs.push_back(vb);
        }
        for (auto b : outputs)
        {
            std::shared_ptr<VkBlob> vb = std::static_pointer_cast<VkBlob, Blob>(b);
            output_blobs.push_back(vb);
        }

        cmd->barrier_writeread(input_blobs, this->top_layer);
        // run inference.
        update_descriptors(input_blobs, output_blobs);
        cmd->bind_pipeline(pipeline);
        cmd->bind_descriptor_set(pipelineLayout, descriptorSet);

        // data size
        pc.input_h = input_h;
        pc.input_w = input_w;
        pc.input_c = input_c;
        pc.output_h = output_h;
        pc.output_w = output_w;
        pc.output_c = output_c;

        cmd->push_constants(pc, pipelineLayout);

        uint32_t nx = (pc.output_w + 31) / 32;
        uint32_t ny = (pc.output_h + 31) / 32;
        cmd->dispatch(nx, ny, 1);
    }
}
